Skip to main content

[SDK] PII Extraction Quickstart (Llama model)

PII Extraction Test with DynamoEval SDK (Llama local model)

Last updated: October 6th, 2024

This Quickstart showcases an end-to-end walkthrough of how to utilize Dynamo AI’s SDK and platform solutions to run privacy tests assessing model vulnerability to PII leakage. For demonstration purposes, we use a GPT-3.5-Turbo model fine-tuned on a customized PolicyQA dataset.

Prerequisites:

  • Dynamo AI API token
  • Hugging Face token (provided by Dynamo AI)
Open In Colab

Environment Setup​

Begin by installing the public Dynamo AI SDK, and importing the required libraries for the quickstart.

DYNAMO_API_KEY=""
DYNAMO_HOST="https://api.dynamo.ai"
# Set your Hugging Face token here -- please use the token provided by DynamoFL to access the Sheared LLaMA-1.3B model for the quickstart
HF_AUTH_TOKEN="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"
!pip install dynamofl==0.0.86

Let's download the model and datasetfiles needed for this quickstart.

!curl "https://www.dropbox.com/scl/fi/seoa6o5tgqxj356uc23q9/quickstart_local_model_pii_extraction.zip?rlkey=088pbmkefy7ea7e1ydrgt30ci&dl=1" -o quickstart_local_model_pii_extraction.zip -J -L -k

!unzip quickstart_local_model_pii_extraction.zip
import time
from dynamofl import DynamoFL, VRAMConfig, GPUConfig, GPUType

Now, create a Dynamo AI instance using your API token and host.

If you do not have an API token, generate a token by logging into apps.dynamo.ai with your provided credentials. Navigate to apps.dynamo.ai/profile to generate your Dynamo AI API token. This API token will enable you to programmatically connect to the Dynamo AI server, create projects, and train models. If you generate multiple API tokens, only your most recent one will work.

dfl = DynamoFL(DYNAMO_API_KEY, host=DYNAMO_HOST)
print(f"Connected as {dfl.get_user()['email']}")

Create a Model​

First, let's create a local model object. The model object specifies the model that privacy tests will be run on during the create_test method. DynamoFL currently supports two types of model objects — local models and remote model API endpoints.

In this quickstart, we demonstrate running tests on local models. A local model object can be used to upload a custom model and run penetration tests.

Creating a local model object requires specifying the model file path and architecture. Currently, DynamoFL supports penetration testing on uploaded models with “.pt” and “.bin” file formats— please confirm that your provided model file fits this formatting. In terms of model architectures, provide any valid HuggingFaceHub model id.

Low-Rank Adaptation (LoRA)​

In this tutorial, we use a local model that has been fine-tuned using Low-Rank Adaptation (LoRA), a technique that that ‘freezes’ the majority of parameters in a pre-trained LLM, while fine-tuning a small number of extra model parameters, leading to a decrease in training time, compute usage, and storage costs. When using a model fine-tuned with LoRA or PEFT (parameter-efficient fine-tuning), it’s necessary to also provide the file path to the peft adapter configuration.

SLUG = time.time()
USE_LORA = True # using a LoRA adapter for a lightweight model
model_dir = "Local Model PII Extraction Quickstart Package/model"

model_file_path = "adapter_model.bin"
model_file_path = os.path.join(model_dir, model_file_path)
peft_config_path = "adapter_config.json"
peft_config_path = os.path.join(model_dir, peft_config_path)
model = dfl.create_model(
  name="Sheared Llama PolicyQA", # model name
  model_file_path=model_file_path,
  architecture="dynamofl-sandbox/sheared-llama-1b3",
  architecture_hf_token=HF_AUTH_TOKEN,
  peft_config_path=peft_config_path
)

print(model.key)

Create a dataset object​

To run a privacy evaluation test, we also need to specify the dataset used for fine-tuning the model. A dataset can be created by specifying the dataset file path. Here, we also provide the dataset with a unique key and an identifying name. At this time, Dynamo AI only accepts csv datasets.

dataset_file_path = "Local Model PII Extraction Quickstart Package/policy_qa_train.csv"  # file path to fine-tuning dataset

dataset = dfl.create_dataset(
    key="dataset_{}".format(SLUG).format(),
    file_path=dataset_file_path,
    name="PolicyQA Finetuning Dataset"
)

print(dataset._id)

Set Test Parameters​

When configuring a test, a variety of test parameters can be configured to customize the test — including the column_name from the dataset to create prompts from, the types of PII to detect leakage for, and the model temperatures to run tests at. Test configuration parameters can be provided to the create_pii_extraction_test method. For more details on best practices for configuring each test parameter, please see the appendix.

PII Entities​

When configuring a PII extraction, inference, or reconstruction attack, one of the most important hyperparameters to configure is the pii_classes parameter. This controls which types of PII the extraction attack will be run for. For the full list of valid PII classes, see the appendix

  • Regex PII Entities In addition to the predefined PII classes, leakage can also be detected for custom-defined regex entities. Custom entities can be added by defining a dictionary mapping entity names to the valid python regex expression in the regex_expressions parameter. Here, we will tag usernames such as john_doe or john_doe123.
pii_entities = ["PERSON", "DATE_TIME", "ORGANIZATION", "EMAIL_ADDRESS"]
regex_expressions = {
		"USERNAME": r"([a-zA-Z]+_[a-zA-Z0-9]+)",
}

Dataset Parameters​

When creating a test, it is required to provide the dataset column names in the test parameters.

When using a decoder-only model (such as GPT), it’s only necessary to provide the pii_ref_column. The identified column will be used for identifying the true PII in the dataset and generating attack prompts.

Hyperparameters​

When running a privacy test, we can also specify a set of hyperparameters to run a sweep of attacks over in the grid parameter. For example, the following grid

grid = [
  {'temperature': [0.1, 0.5, 1.0], 'seq_len': [64, 128]},
  {'temperature': [0.5, 1.0], 'seq_len': [256]},
 ]

specifies that two sets of attacks should be run: one with temperature values in [0.1, 0.5, 1.0] and sequence length values in [64, 128], and the second one with temperature values in [0.5, 1.0] and a sequence length of 256. For PII extraction tests, the following hyperparameters can be explored:

  • temperature: Temperature controls the amount of ‘randomness’ in your models generations. Varying temperature may significantly affect attack results.

  • seq_len: The sequence length parameter is available for only PII Extraction attacks and controls the number of tokens to generate per attack. Increasing the sequence length will result in longer model responses, which may increase the probability of PII being emitted. We recommend setting the sequence length based on how your model may be used in production.

Finally, for PII extraction tests, we can also specify the sampling_rate parameter. The sampling rate parameter is also available for only PII Extraction attacks and controls the number of times the model is prompted during the attack. Increasing the sampling rate will result in a more expensive, but more robust test. Dynamo AI recommends a default of 1024.

Run tests on LLaMA2-1.3B​

To run a PII extraction privacy evaluation test, we can call the create_pii_extraction_test method. Test creation will submit a test to our DynamoEval platform, where the test will be run.

test_info = dfl.create_pii_extraction_test(
  name="privacy_test_{}".format(SLUG).format(),
  model_key=model.key,
  dataset_id=dataset._id,
  gpu=GPUConfig(gpu_type=GPUType.A10G, gpu_count=1),
  pii_ref_column="text",
  pii_classes=pii_entities,
  sampling_rate=1024,
  regex_expressions=regex_expressions,
  grid=[
    {'temperature': [0.1, 0.5, 1.0]},
  ]
)

Checking the attack status​

# Confirming the Attack has been queued
def query_attack_status(attack_id):
    attack_info = dfl.get_attack_info(attack_id)['status']
    print("Attack status: {}.".format(attack_info))

all_attacks = test_info.attacks
attack_ids = [attack["id"] for attack in all_attacks]
for attack_id in attack_ids:
  print(query_attack_status(attack_id))

Viewing test results​

After your test has been created, navigate to the model dashboard page in the Dynamo AI UI. Here, you should observe that your model has been created and that your test is running.
After the test has completed, a test report will be created and you can dive into the test results!